import os
from typing import Tuple, List, Any
import chardet

# 得到训练数据，第一个返回值是所有svs文件的特征向量，第二个返回值是所有svs对应的label(是否转移)
import numpy as np
import pandas as pd
import torch
from numpy import ndarray
from sklearn.model_selection import train_test_split

# MLP模型定义
from torch import nn, optim
from torch.utils.data import DataLoader, TensorDataset


# MLP（Multi-Layer Perceptron，多层感知器）是一种基础的人工神经网络模型，广泛用于监督学习任务，尤其是在分类和回归问题中。
# MLP 的优势：
# 简单易用：结构简单，适用于各种问题，尤其是中小型数据集。
# 灵活性强：适用于分类、回归等任务。
#
# MLP 的缺点：
# 过拟合风险高：在数据量较小或特征复杂的情况下，容易出现过拟合。
# 缺乏结构信息：相比于卷积神经网络（CNN）等，MLP 对图像、语音等结构化数据的处理效果差。
class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(128 * 128, 512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, 128),
            nn.ReLU(),
            nn.Linear(128, 2)
        )

    def forward(self, x):
        return self.model(x)


# 第一个ndarray所有svs数据的特征向量，第二个ndarray是所有svs的标签
def get_train_data(label_csv: str, svs_feature_path: str) -> Tuple[ndarray, ndarray]:
    # 用 Excel 打开 label.csv，点击「文件」>「另存为」，保存类型选：CSV UTF-8（逗号分隔）(*.csv)
    label_df = pd.read_csv(label_csv, encoding='utf-8')  # 假设有列：slide_id, label
    # 加载特征
    svs_feature = []
    svs_label = []
    for _, row in label_df.iterrows():
        svs_name = row["slide_id"]
        label = row["label"]
        feat_path = os.path.join(svs_feature_path, f"{svs_name}.csv")  # 根据实际路径改
        feat = np.loadtxt(feat_path, delimiter=",")  # shape: [128, 128]
        svs_feature.append(feat)
        svs_label.append(label)
    # shape: [svs数量, 128, 128] 表示svs数量与每个svs的表示特征
    X = np.array(svs_feature)
    # shape: [svs数量] 表示标签
    Y = np.array(svs_label)
    return X, Y


# 超参数：
# epochs：整个训练集会被“完整跑一遍”的次数。太小 → 模型学不够，欠拟合，太大 → 模型可能学过头，过拟合。
# batch_size：每次训练时“喂给模型的样本数量”。小 → 更稳定但训练慢（适合小数据或内存少）；大 → 更快但可能震荡或内存爆炸。batch_size 通常选择为 4、8、16、32、64、128 等 2 的幂。  例子：如果有64个数据，batch_size=8，每次训练需要喂8次数据
# lr: 学习率，每次参数更新的“步子”大小。太小 → 学得慢，可能卡住；太大 → 学得快但可能错过最佳解（震荡甚至发散）。
def train_svs_classifier(feature_dir: str, label_csv: str,
                         epochs: int = 30, batch_size: int = 8,
                         lr: float = 1e-4, model_save_path: str = "svs_model.pth"):
    # 第一步：获取数据
    data = get_train_data(label_csv, feature_dir)
    print("所有svs特征表示: ", data[0].shape)
    print("所有svs特征的label: ", data[1].shape, data[1])

    # 第二步：划分训练集与验证集
    # test_size=0.2：表示 80% 的数据用于训练，20% 用于验证。
    # stratify=y：按类别标签比例进行分层抽样，例如如果你的标签是 1 和 0，且原始数据中 1 和 0 各占一半，那么划分后，训练集和验证集中也会维持这种比例，防止数据不均衡带来的偏差。
    # random_state=42：设置随机种子，使结果可复现（只要数据不变，每次划分都一样）。
    X_train, X_val, y_train, y_val = train_test_split(data[0], data[1], test_size=0.3, stratify=data[1],
                                                      random_state=42)
    print(f"训练集样本数: {len(X_train)}")
    print(f"验证集样本数: {len(X_val)}")
    print(f"训练集正负样本分布: {np.bincount(y_train)}")
    print(f"验证集正负样本分布: {np.bincount(y_val)}")

    # 第三步： 训练集与验证集转为 Tensor，是将训练集与测试集转为可训练的数据
    X_train_tensor = torch.tensor(X_train, dtype=torch.float32).view(-1, 128 * 128)
    X_val_tensor = torch.tensor(X_val, dtype=torch.float32).view(-1, 128 * 128)
    y_train_tensor = torch.tensor(y_train, dtype=torch.long)
    y_val_tensor = torch.tensor(y_val, dtype=torch.long)

    train_loader = DataLoader(TensorDataset(X_train_tensor, y_train_tensor), batch_size=batch_size, shuffle=True)
    val_loader = DataLoader(TensorDataset(X_val_tensor, y_val_tensor), batch_size=batch_size)

    # 第四步：模型与优化器，是定义训练模型。
    # 定义训练所需的核心三要素：
    # model：用 MLP() 定义了一个多层感知机分类模型；
    # optimizer：使用 Adam 优化器来更新模型权重；
    # criterion：使用交叉熵损失函数，用于分类问题衡量预测与真实标签的差距。
    # PyTorch 的训练流程：口诀 --->>> "数据张量化，模型定义好，损失 + 优化器，循环跑训练"
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")
    model = MLP().to(device)
    optimizer = optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()

    # 第五步： 开始训练
    for epoch in range(epochs):
        model.train()
        total_loss = 0
        for xb, yb in train_loader:
            xb, yb = xb.to(device), yb.to(device)
            optimizer.zero_grad()
            preds = model(xb)
            loss = criterion(preds, yb)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        # 验证
        model.eval()
        correct = 0
        total = 0
        with torch.no_grad():
            for xb, yb in val_loader:
                xb, yb = xb.to(device), yb.to(device)
                preds = model(xb)
                pred_labels = preds.argmax(dim=1)
                correct += (pred_labels == yb).sum().item()
                total += yb.size(0)
        # Loss：表示模型在训练集上的损失（误差），通常损失越低表示模型的表现越好。初始的损失值可能比较高，因为模型还没有学到有效的特征。
        # Val Acc：表示模型在验证集上的准确率，0.5 意味着模型的预测结果和实际标签的一致性为 50%。
        print(f"Epoch {epoch + 1}: Loss = {total_loss:.4f}, Val Acc = {correct / total:.4f}")

    # 保存模型
    torch.save(model.state_dict(), model_save_path)
    print("模型训练完成！！！！")
    print(f"✅ 模型已保存到 {model_save_path}")


if __name__ == '__main__':
    train_svs_classifier(
        feature_dir="D:/zm_scientific_research/svs_features",  # 存放每个svs的128x128特征csv
        label_csv="D:/zm_scientific_research/transfer/label.csv",
        model_save_path="D:/zm_scientific_research/transfer/svs_transfer_model.pth",
        epochs=60,
        batch_size=2
    )
